using namespace std;
#include<iostream>
#include<iomanip>
#include<fstream>
#include<vector>
#include<stdio.h>
#include<stdlib.h>
#include<map>
#include<set>
#include"BA2.h"
#include<math.h>
#include<gsl/gsl_rng.h>
#include<gsl/gsl_randist.h>
#include<gsl/gsl_sf_gamma.h>

extern std::vector<string> cmdList;
extern std::vector<Indiv> sample;
extern std::map<string,int> popToindex;
extern std::map<int,string> indexTopop;
extern std::map<string,int> labelToindex;
extern std::map<int,string> indexTolabel;
extern std::set<string> labelSet;
extern std::set<string> popSet;
extern std::vector<long int> popCounts;
extern std::vector< vector< vector<long int> > > alleleCounts;
extern std::vector< set<string> > alleleSet;
extern std::vector< map<string,int> > alleleMap;
extern std::vector< map<int,string> > intToallele;
extern std::set<string> locusSet;
extern std::map<int,string> indexTolocus;
extern std::map<string,int> locusToindex;
extern double maxMigration;
extern bool infileOpen;

// constructor for Parameters class
Parameters::Parameters()
{
  logfile="BA2.log";
  outfile="BA2.out";
  Ngens=0;
  MCMCreps=50000;
  MCMCthin=1000;
  MCMCdisplay=1000;
  MCMCburnin=5000;
  migPropDirichlet=100;
  alleleFreqDirichlet=100;
  windowDelta=0.1;
  missingData="?";
  seed=1;
}

int TranslateCommand(string cmd)
{
  // parse command and options into vector of strings
  cmdList.erase(cmdList.begin(),cmdList.end());

  const string delims(" \t");
  string::size_type begIdx, endIdx;
  begIdx=cmd.find_first_not_of(delims);
  while(begIdx!=string::npos)
    {
      endIdx = cmd.find_first_of(delims,begIdx);
      if(endIdx==string::npos)
	{
	  endIdx=cmd.length();
	}
      cmdList.push_back(cmd.substr(begIdx,endIdx-begIdx));
      begIdx=cmd.find_first_not_of(delims,endIdx);
    }

  // translate command into integer
  if(cmdList[0]=="quit") return(1);
  else
    if(cmdList[0]=="scan") return(2);
    else
      if(cmdList[0]=="pset") return(3);
      else
	if(cmdList[0]=="help") return(4);
	else
	  if(cmdList[0]=="printData") return(5);
	  else
	    if(cmdList[0]=="printAlleles") return(6);
	    else
	      if(cmdList[0]=="bat") return(7);
	      else
		if(cmdList[0]=="bayesAss") return(8);
		else
		  if(cmdList[0]=="simulate") return(9);
		  else
		    if(cmdList[0]=="sset") return(10);
		    else
		      if(cmdList[0]=="interact") return(12);
		      else
			return(0);
}

void PrintBanner()
{
  cout << "             BayesAss V.2.0 [FlatHead Lake Edition]              " << "\n";
  cout << "               Bruce Rannala (http://rannala.org)                " << "\n";     
  cout << "         Genome Center, University of California, Davis          " <<  "\n";
  cout << "            Copyright 2007, University of California             " << "\n";
  cout << "               [enter \"help;\" to list commands]                " << "\n";
}

void PrintCommands()
{
	  cout << "                                                            \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "               BayesAss V.2.0: List of Commands             \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " bat [-I include/exclude] : Bayesian population assignment  \n"; 
	  cout << " bayesAss : Bayesian MCMC analysis of migrant proportions   \n"; 
	  cout << " citation [-C citation format] : articles to cite for BA2   \n";
	  cout << " help : print list of commands to screen                    \n";
	  cout << " printAlleles : print allele counts                         \n";
	  cout << " printData : print to screen data scanned from input file   \n";	
	  cout << " pset [parameter=value] : set parameter to value            \n";
	  cout << " scan [-F file format][input file name] : open input file   \n";
	  cout << " simulate : simulate data                                   \n";
	  cout << " sset [parameter=value] : set simulation parameter to value \n";
	  cout << " quit : terminate the program                               \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " Important note: all commands must be terminated with \";\" \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "                                                            \n";
}

void DataSummary(Parameters pvalues)
{
  // print summary to screen
  cout << "Data File in " << pvalues.fileFormat << " format \n";
  cout << "Individuals: " << sample.size() << " ";
  cout << "Populations: " << popSet.size() << " ";
  cout << "Loci: " << sample[0].genotypes.size() << "\n";
}

void ReadData(ifstream &infile, Parameters &pvalues)
{
  string aLine;
  Gene geno;
  Indiv temp;

  cout << "Scanning input file: " << cmdList[2] << "\n";
  sample.erase(sample.begin(),sample.end());
  popSet.erase(popSet.begin(),popSet.end());
  popCounts.erase(popCounts.begin(),popCounts.end());
  popToindex.erase(popToindex.begin(),popToindex.end());
  indexTopop.erase(indexTopop.begin(),indexTopop.end());
  labelToindex.erase(labelToindex.begin(),labelToindex.end());
  indexTolabel.erase(indexTolabel.begin(),indexTolabel.end());
  alleleCounts.erase(alleleCounts.begin(),alleleCounts.end());
  alleleSet.erase(alleleSet.begin(),alleleSet.end());
  alleleMap.erase(alleleMap.begin(),alleleMap.end());
  intToallele.erase(intToallele.begin(),intToallele.end());
  labelSet.erase(labelSet.begin(),labelSet.end());
  pvalues.alist.erase(pvalues.alist.begin(),pvalues.alist.end());
  
  int i=0, popindex=0, indivindex=0, locusindex=0;

  if(pvalues.fileFormat=="-FS")      // data is in Structure file format
    {
      while(getline(infile,aLine))
	{
	  i++;
	  vector<string> parsedLine;
	  const string delims(" \t");
	  string::size_type begIdx, endIdx;
	  begIdx=aLine.find_first_not_of(delims);
	  while(begIdx!=string::npos)
	    {
	      endIdx = aLine.find_first_of(delims,begIdx);
	      if(endIdx==string::npos)
		{
		  endIdx=aLine.length();
		}
	      parsedLine.push_back(aLine.substr(begIdx,endIdx-begIdx));
	      begIdx=aLine.find_first_not_of(delims,endIdx);
	    }
	  if((i % 2)!=0)
	    {
	      temp.label=parsedLine[0];
	      labelSet.insert(parsedLine[0]);
	      labelToindex[string(parsedLine[0])]=indivindex;
	      indexTolabel[indivindex]=parsedLine[0];
	      ++indivindex;
	      temp.sampledPopulation=parsedLine[1];
	      if(popSet.count(parsedLine[1])==0)
		{
		  popSet.insert(parsedLine[1]);
		  popToindex[string(parsedLine[1])]=popindex;
		  indexTopop[popindex]=parsedLine[1];
		  ++popindex;
		}
	      for(int j=2;j<=static_cast<int>(parsedLine.size())-1;j++)
		{ 
		  geno.g1=parsedLine[j];
		  temp.genotypes.push_back(geno);
		} 
	    }
	  else
	    {
	      for(int j=2;j<=static_cast<int>(parsedLine.size())-1;j++)
		{ 
		  temp.genotypes[j-2].g2=parsedLine[j];
		} 
	      sample.push_back(temp);
	      temp.genotypes.erase(temp.genotypes.begin(),temp.genotypes.end());
	    }
	} 
    }
  else
    if(pvalues.fileFormat=="-FI")      // data is in Immanc file format
      {
	vector<vector <string> > parsedLines;
	// put all lines from input file into parsedLines
	while(getline(infile,aLine))
	  {
	    vector<string> oneParsedLine;
	    const string delims(" \t");
	    string::size_type begIdx, endIdx;
	    begIdx=aLine.find_first_not_of(delims);
	    while(begIdx!=string::npos)
	      {
		endIdx = aLine.find_first_of(delims,begIdx);
		if(endIdx==string::npos)
		  {
		    endIdx=aLine.length();
		  }
		oneParsedLine.push_back(aLine.substr(begIdx,endIdx-begIdx));
		begIdx=aLine.find_first_not_of(delims,endIdx);
	      }
	    if(static_cast<int>(oneParsedLine.size())>1)
	      parsedLines.push_back(oneParsedLine);
	  }

	// get set of labels, locus, and population names and create index maps
	for(int j=0;j<static_cast<int>(parsedLines.size());j++)
	  {
	    if(labelSet.count(parsedLines[j][0])==0)
	      {
		labelSet.insert(parsedLines[j][0]);
		labelToindex[string(parsedLines[j][0])]=indivindex;
		indexTolabel[indivindex]=parsedLines[j][0];
		++indivindex;
	      }
	    if(popSet.count(parsedLines[j][1])==0)
	      {
		popSet.insert(parsedLines[j][1]);
		popToindex[string(parsedLines[j][1])]=popindex;
		indexTopop[popindex]=parsedLines[j][1];
		++popindex;
	      }
	    if(locusSet.count(parsedLines[j][2])==0)
	      {
		locusSet.insert(parsedLines[j][2]);
		locusToindex[string(parsedLines[j][2])]=locusindex;
		indexTolocus[locusindex]=parsedLines[j][2];
		++locusindex;
	      }
	  }
	
	geno.g1=""; 
	geno.g2="";
	for(int j=0;j<static_cast<int>(locusSet.size());j++)
	  temp.genotypes.push_back(geno);
	for(int j=0;j<static_cast<int>(labelSet.size());j++)
	  {
	    temp.label=indexTolabel[j];
	    int k=0; 
	    bool notFoundPop=true;
	    while((k<static_cast<int>(parsedLines.size()))&&notFoundPop)
	      {
		if(labelToindex[parsedLines[k][0]]==j)
		  { temp.sampledPopulation=parsedLines[k][1]; notFoundPop=false; }
		k++;
	      }
	    for(k=0;k<static_cast<int>(parsedLines.size());k++)
	      {
		if(labelToindex[parsedLines[k][0]]==j)
		  {
		    temp.genotypes[locusToindex[parsedLines[k][2]]].g1=parsedLines[k][3];
		    temp.genotypes[locusToindex[parsedLines[k][2]]].g2=parsedLines[k][4];
		  }
	      }
	    sample.push_back(temp);
	  }
      }
    else
      {
	cout << "error: file format " << pvalues.fileFormat << " is not recognized\n";
	infileOpen=false;
      }

  if(infileOpen)
    {
      // Get counts of alleles sampled from each population
      GetAlleleCounts(pvalues);

      // Get counts of individuals sampled from each population
      for(int i=0;i<static_cast<int>(popSet.size());i++)
	popCounts.push_back(0);
      for(int i=0;i<static_cast<int>(popSet.size());i++)
	for(int j=0;j<static_cast<int>(sample.size());j++)
	  if(popToindex[sample[j].sampledPopulation]==i)
	    popCounts[i]+=1;
      
      // print data summary to screen
      DataSummary(pvalues);
    }
}

void PrintData(string filename)
{
  string aString;
 
  cout << "\nPrinting Contents of Data File: " << filename << "\n";
  cout << "------------------------------------------------------------\n";
  cout << " Indiv.    Popln.                   Genotypes\n";
  cout << "------------------------------------------------------------\n";
  cout << setfill(' ');
  for(int j=0;j<static_cast<int>(sample.size());j++)
    {
      cout << " " << setw(10) << setiosflags(ios::left) << sample[j].label;
      cout << sample[j].sampledPopulation << "         ";
      for(int l=0;l<static_cast<int>(sample[j].genotypes.size());l++)
	{
	  aString=sample[j].genotypes[l].g1+"/"+sample[j].genotypes[l].g2;
	  cout <<  setw(10) << setiosflags(ios::left) << aString;
	}
      cout << "\n";
    }
  cout << "------------------------------------------------------------\n";
  cout << "\n";
}

void PrintParameters(Parameters pvalues)
{
	  cout << "                                                            \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "               BayesAss V.2.0: Parameter Settings           \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " alist=(";
	  if(pvalues.alist.size()==0)
	    cout << ")\n";
	  else
	    for(int i=0;i<=static_cast<int>(pvalues.alist.size())-1;i++)
	      {
		cout << pvalues.alist[i];
		if(i!=(static_cast<int>(pvalues.alist.size())-1)) cout << ",";
		else cout << ")\n";
	      }
	  cout << " MCMCreps=" << pvalues.MCMCreps << "\n";
	  cout << " MCMCthin=" << pvalues.MCMCthin << "\n";
	  cout << " MCMCdisplay=" << pvalues.MCMCdisplay << "\n";
	  cout << " MCMCburnin=" << pvalues.MCMCburnin << "\n";
	  cout << " MCMCfstatPr=" << pvalues.windowDelta << "\n";
	  cout << " MCMCfreqPr=" << pvalues.alleleFreqDirichlet << "\n";
	  cout << " MCMCmigrPr=" << pvalues.migPropDirichlet << "\n";
	  cout << " Ngens=" << pvalues.Ngens << "\n";
	  cout << " missingData=" << pvalues.missingData << "\n";
	  cout << " outfile=" << pvalues.outfile << "\n";
	  cout << " logfile=" << pvalues.logfile << "\n";
	  cout << " RandomSeed=" << pvalues.seed << "\n";
	  cout << "------------------------------------------------------------\n";
	  cout << "                                                            \n";
}

int SetParameters(Parameters &pvalues)
{
  string parm;
  string paramValue;
  int file_changed=0;
  if(cmdList.size()==1)
    {
      PrintParameters(pvalues);
    }
  else
    if(static_cast<int>(cmdList.size())==2)
      {
	int pos=cmdList[1].find_first_of("=");
	if(pos!=static_cast<int>(string::npos))
	  {
	    parm=cmdList[1].substr(0,pos);
	    paramValue=cmdList[1].substr(pos+1);
	  }
	if(parm=="MCMCreps") pvalues.MCMCreps=atol(paramValue.c_str());
	else
	  if(parm=="MCMCthin") pvalues.MCMCthin=atol(paramValue.c_str());
	  else
	    if(parm=="MCMCdisplay") pvalues.MCMCdisplay=atol(paramValue.c_str());
	    else
	      if(parm=="MCMCburnin") pvalues.MCMCburnin=atol(paramValue.c_str());
	      else
		if(parm=="outfile") {pvalues.outfile=paramValue; file_changed=1;}
		else 
		  if(parm=="logfile") {pvalues.logfile=paramValue; file_changed=2;}
		  else
		    if(parm=="Ngens") 
		      {
			int temp=atol(paramValue.c_str());
			if((temp==0)||(temp==1))
			  pvalues.Ngens=temp;
			else
			  cout << "Error: Ngens=" << temp << " (should be either 0 or 1)\n";
		      }
		  else
		    if(parm=="MCMCfreqPr") pvalues.alleleFreqDirichlet=atof(paramValue.c_str());
		    else
		      if(parm=="MCMCmigrPr") pvalues.migPropDirichlet=atof(paramValue.c_str());
		      else
			if(parm=="MCMCfstatPr") pvalues.windowDelta=atof(paramValue.c_str());
			else
			  if(parm=="missingData") pvalues.missingData=paramValue;
			  else
			    if(parm=="RandomSeed") pvalues.seed=atol(paramValue.c_str());
			    else
			      if(parm=="alist") 
				{ 
				  vector<string> tempList;
				  paramValue.erase(paramValue.begin(),paramValue.begin()+1);
				  int pos=0; int prev_pos=0;
				  while((pos=paramValue.find_first_of(',',pos))!=static_cast<int>(string::npos))
				    {
				      tempList.push_back(paramValue.substr(prev_pos,pos-prev_pos));
				      prev_pos = ++pos;
				    }
				  tempList.push_back(paramValue.substr(prev_pos,pos-prev_pos));
				  tempList[static_cast<int>(tempList.size())-1]= 
				    tempList[tempList.size()-1].erase(static_cast<int>(tempList[static_cast<int>(tempList.size())-1].size())-1,static_cast<int>(tempList[static_cast<int>(tempList.size())-1].size()));
				  
				  // check that individuals exist in sample, otherwise return error
				  bool missingLabel=0;
				  for(int i=0;i<=static_cast<int>(tempList.size())-1;i++)
				    if(labelSet.count(tempList[i])==0) {cout << "Error: label " << tempList[i] << " not in sample\n"; missingLabel=1;} 
				  if(!missingLabel)
				    pvalues.alist=tempList;
				}
      }
  PrintParameters(pvalues);
  return(file_changed);
}

void GetAlleleCounts(Parameters pvalues)
{
  // create map from allele labels to index numbers
  for(int i=0;i<static_cast<int>(sample[0].genotypes.size());i++)
    {
      set<string> tempSet;
      map<string,int> tempMap;
      map<int,string> tempintToallele;
      int alleleit=1;
      // always insert missing data character into allele set 
      tempSet.insert(pvalues.missingData);
      tempMap[pvalues.missingData]=0;
      tempintToallele[0]=pvalues.missingData;
      for(int j=0;j<=static_cast<int>(sample.size())-1;j++)
	{
	  if(tempSet.count(sample[j].genotypes[i].g1)==0)
	    {
	      tempSet.insert(sample[j].genotypes[i].g1);
	      tempMap[string(sample[j].genotypes[i].g1)]=alleleit;
	      tempintToallele[alleleit]=sample[j].genotypes[i].g1;
	      ++alleleit;
	    }
	  if(tempSet.count(sample[j].genotypes[i].g2)==0)
	    {
	      tempSet.insert(sample[j].genotypes[i].g2);
	      tempMap[string(sample[j].genotypes[i].g2)]=alleleit;
	      tempintToallele[alleleit]=sample[j].genotypes[i].g2;
	      ++alleleit;
	    }
	}
      alleleSet.push_back(tempSet);
      alleleMap.push_back(tempMap);
      intToallele.push_back(tempintToallele);
    }

  // create data structure for allele counts
  vector< vector<long int> > tempPopln;
  for(int j=0;j<=static_cast<int>(sample[0].genotypes.size())-1;j++)
    {
      vector<long int> tempLocus;
      for(int k=0; k<=static_cast<int>(alleleSet[j].size())-1; k++)
	tempLocus.push_back(0);
      tempPopln.push_back(tempLocus);
    }
  for(int i=0;i<=static_cast<int>(popSet.size())-1;i++)
    {
      alleleCounts.push_back(tempPopln);
    }

  // tabulate allele counts
  for(int i=0;i<=static_cast<int>(sample.size())-1;i++)
    for(int j=0;j<=static_cast<int>(sample[0].genotypes.size())-1;j++)
      {
	alleleCounts[popToindex[string(sample[i].sampledPopulation)]][j][alleleMap[j][string(sample[i].genotypes[j].g1)]]++;
	alleleCounts[popToindex[string(sample[i].sampledPopulation)]][j][alleleMap[j][string(sample[i].genotypes[j].g2)]]++;
      }
}

void PrintAlleleCounts(Parameters pvalues)
{
  cout << "\nMissing Data = " << pvalues.missingData << "\n";
  for(int i=0;i<=static_cast<int>(popSet.size())-1;i++)
    {
      cout << "\nPopulation: " << indexTopop[i] << "\n";
      for(int j=0;j<=static_cast<int>(alleleCounts[i].size())-1;j++)
	{
	  cout << "Locus " << j << ": ";
	  for(int k=0;k<=static_cast<int>(alleleCounts[i][j].size())-1;k++)
	    cout << intToallele[j][k]<< ":" << alleleCounts[i][j][k] << " ";
	  cout << "\n";
	}
    } 
  cout << "\n";
}

double logGenotypeProbability(int indiv, int popln, vector< vector< vector<long int> > > ExAlleleCounts)
{
  double logPr=0.0;
  for(int i=0;i<static_cast<int>(sample[0].genotypes.size());i++)
    {
      if(sample[indiv].genotypes[i].g1==sample[indiv].genotypes[i].g2)
	{
	  logPr+=log(ExAlleleCounts[popln][i][alleleMap[i][sample[indiv].genotypes[i].g1]]+
		     1.0/static_cast<int>(ExAlleleCounts[popln][i].size())+1.0)+
	    log(ExAlleleCounts[popln][i][alleleMap[i][sample[indiv].genotypes[i].g1]]+
		1.0/static_cast<int>(ExAlleleCounts[popln][i].size()));
	  int totIndiv=0;
	  for(int j=0;j<static_cast<int>(ExAlleleCounts[popln][i].size());j++)
	    totIndiv+=ExAlleleCounts[popln][i][j];
	  logPr-=log(totIndiv+2);
	  logPr-=log(totIndiv+1);
	}
      else
	{
	  logPr+=log(alleleCounts[popln][i][alleleMap[i][sample[indiv].genotypes[i].g1]]
		     +1.0/static_cast<int>(alleleCounts[popln][i].size()))+
	    log(alleleCounts[popln][i][alleleMap[i][sample[indiv].genotypes[i].g2]]+
		1.0/static_cast<int>(alleleCounts[popln][i].size()));
	  int totIndiv=0;
	  for(int j=0;j<static_cast<int>(alleleCounts[popln][i].size());j++)
	    totIndiv+=alleleCounts[popln][i][j];
	  logPr-=log(totIndiv+2);
	  logPr-=log(totIndiv+1);
	}
    }
  return(logPr);
}

double BayesianAssignment(int indiv, int popln,	vector< vector< vector<long int> > > ExAlleleCounts)
{
  double sumProb=0.0;
  for(int j=0;j<static_cast<int>(alleleCounts.size());j++)
    sumProb+=exp(logGenotypeProbability(indiv,j,ExAlleleCounts));
  return(exp(logGenotypeProbability(indiv,popln,ExAlleleCounts))/sumProb);
}

void Bat(Parameters pvalues)
{
  bool exclude=true;
  cout << "------------------------------------------------------------\n";
  cout << "Bayesian Population Assignment\n";
  cout << "------------------------------------------------------------\n";
  cout << setfill(' ');
  if(cmdList.size()>1)
    if(cmdList[1]=="-I") 
      { 
	exclude=false; cout << "--Assigned individual included to calculate probability--\n";
      }
  cout << "    Population:  ";
  for(int i=0;i<static_cast<int>(alleleCounts.size());i++)
    cout << setw(10) <<  setiosflags(ios::left) << indexTopop[i];
  cout << "\n";
  cout << "Individual:";
  for(int j=0;j<static_cast<int>(pvalues.alist.size());j++)
    {
      std::vector< vector< vector<long int> > > ExAlleleCounts=alleleCounts;
      if(exclude==true) // remove allele counts for individual being assigned 
	{
	  for(int k=0;k<static_cast<int>(sample[0].genotypes.size());k++)
	    {
	      --ExAlleleCounts[popToindex[sample[labelToindex[pvalues.alist[j]]].sampledPopulation]][k][alleleMap[k][sample[labelToindex[pvalues.alist[j]]].genotypes[k].g1]];
	      --ExAlleleCounts[popToindex[sample[labelToindex[pvalues.alist[j]]].sampledPopulation]][k][alleleMap[k][sample[labelToindex[pvalues.alist[j]]].genotypes[k].g2]];
	    }
	}  	
      cout << "\n" << setw(17) << pvalues.alist[j];
      for(int i=0;i<static_cast<int>(alleleCounts.size());i++)
	{
	  cout << setw(10) << setprecision(3) << setiosflags(ios::left) << BayesianAssignment(labelToindex[pvalues.alist[j]],i,ExAlleleCounts);
	}
    }
  cout << "\n------------------------------------------------------------\n\n";
}

